import warnings
warnings.filterwarnings('ignore')
import os
import cv2

import numpy as np
import torch
import onnx
from onnxsim import simplify
from torch import nn
from torch.nn import functional as F
class ConvNormAct(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=(0, 0), dilation=(1, 1), bias=False):
        super(ConvNormAct, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=padding, stride=stride,
                      bias=bias,dilation=dilation),
            nn.BatchNorm2d(out_channels, momentum=0.1),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        x = self.conv(x)
        return x
class ATT48(nn.Module):
    def __init__(self, inplanes):
        super(ATT48, self).__init__()
        midplanes = inplanes//4
        self.pool1 = nn.AvgPool2d(kernel_size=4, stride=4)
        self.pool2 = nn.AvgPool2d(kernel_size=8, stride=8)
        self.pool3 = nn.MaxPool2d(kernel_size=4, stride=4)
        self.pool4 =  nn.MaxPool2d(kernel_size=8, stride=8)
        self.ip = inplanes//4
        self.act = nn.ReLU(inplace=True)
        self.conv_1 = nn.Sequential(
            nn.Conv2d(midplanes, midplanes, kernel_size=1, padding=0, bias=False),
            nn.BatchNorm2d(midplanes),
        )
        self.conv_2 = nn.Sequential(
            nn.Conv2d(midplanes, midplanes, kernel_size=1, padding=0, bias=False),
            nn.BatchNorm2d(midplanes),
        )
        self.conv_3 = nn.Sequential(
            nn.Conv2d(midplanes, midplanes, kernel_size=1, padding=0, bias=False),
            nn.BatchNorm2d(midplanes),
        )
        self.conv_4 = nn.Sequential(
            nn.Conv2d(midplanes, midplanes, kernel_size=1, padding=0, bias=False),
            nn.BatchNorm2d(midplanes),

        )

    def forward(self, x):
        _, _, h, w = x.size()
        x1,x2,x3, x4 = torch.split(x, self.ip, dim=1)
        x1 = F.interpolate(self.conv_1(self.pool1(x1)), scale_factor=4, mode='bilinear')
        x2 = F.interpolate(self.conv_2(self.pool2(x2)), scale_factor=8, mode='bilinear')
        x3 = F.interpolate(self.conv_3(self.pool3(x3)),  scale_factor=4, mode='bilinear')
        x4 = F.interpolate(self.conv_4(self.pool4(x4)),  scale_factor=8, mode='bilinear')
        return torch.sigmoid(torch.cat([x1,x2,x3,x4],dim=1))*x
class CovBlock48(nn.Module):
    expansion = 1
    def __init__(self, inplanes):
        super(CovBlock48, self).__init__()
        self.conv1 = ConvNormAct(inplanes, inplanes, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(inplanes, inplanes, kernel_size=3, stride=1, padding=1, bias=False)
        self.att = ATT48(inplanes)
        self.bn2 = nn.BatchNorm2d(inplanes , momentum=0.1)
        self.relu = nn.ReLU(inplace=True)
    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.att(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out += residual
        out = self.relu(out)
        return out
class R2DliaBlock(nn.Module):
    expansion = 1
    def __init__(self, inplanes, planes, scale_factor=1):
        super(R2DliaBlock, self).__init__()
        self.relu = nn.ReLU(inplace=True)
        self.width = inplanes*scale_factor
        self.ip = self.width//4
        self.conv0 = ConvNormAct(inplanes, self.width, kernel_size=3, stride=1, padding=1)
        self.conv1 = ConvNormAct(self.ip, self.ip, kernel_size=1, padding=0)
        self.conv2 = ConvNormAct(self.ip, self.ip, kernel_size=3, padding=1)
        self.conv3 = ConvNormAct(self.ip, self.ip, kernel_size=3, padding=2, dilation=2)
        self.conv4 = ConvNormAct(self.ip, self.ip, kernel_size=3, padding=3, dilation=3)
        self.conv_out = nn.Sequential(
            nn.Conv2d(self.width, planes, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(planes),
        )
    def forward(self, x):
        residual = x
        out = self.conv0(x)
        x1, x2, x3, x4 = torch.split(out, self.ip, dim=1)
        x1 = self.conv1(x1)
        x2 = self.conv2(x1+x2)
        x3 = self.conv3(x2+x3)
        x4 = self.conv4(x3+x4)
        out = torch.cat([x1, x2, x3, x4], dim=1)
        out = self.conv_out(out)
        out += residual
        out = self.relu(out)
        return out
#model = TWoStageSSD(num_class=num_class, fpn_nums=3,scores_the=0.4, top_detect=2000)
model = R2DliaBlock(128,128)
#model.load_state_dict(torch.load('/data/check/edge_model_1219_544_456.pth'),strict=True)
#model.load_state_dict(torch.load('/home/dsl/release/xag_fpv_1207_scale2267.pth'),strict=True)


model.eval()
img = torch.randn(1,128, 64, 64, requires_grad=True)
torch_out = model(img)
torch.onnx.export(model,               # model being run
                  img,           # model input (or a tuple for multiple inputs)
                  "block.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=11,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  verbose=False,
                 )

onnx_model = onnx.load("block.onnx")
onnx.checker.check_model(onnx_model)
onnx_model = onnx.load('block.onnx')
sim, _ =simplify(onnx_model)
onnx.save_model(sim, 'block.onnx')